package de.unibi.comet.ac;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;

import de.unibi.comet.fa.Alphabet;
import de.unibi.comet.fa.CDFA;
import de.unibi.comet.hmm.Emission;
import de.unibi.comet.hmm.HMM;
import de.unibi.comet.hmm.Node;
import de.unibi.comet.hmm.StandardNode;
import de.unibi.comet.hmm.Transition;
import de.unibi.comet.util.Log;
import de.unibi.comet.util.PreOrderIterator;

public class AhoCorasickAutomaton {
	private AhoCorasickNode root;

	public AhoCorasickAutomaton() {
		clear();
	}

	public void clear() {
		root = null;
	}

	public String pstricksOutput() {
		if (root == null)
			return "";
		return root.pstricksOutput();
	}

	private class FailureLinkFinder {
		private AhoCorasickNode[] nodeTable;

		
		FailureLinkFinder(AhoCorasickNode[] nodeTable) {
			this.nodeTable = nodeTable;
		}

		/**
		 * Finds trie node which corresponds to suffix tree node of reversed
		 * joined string.
		 */
		AhoCorasickNode translateNode(SuffixStructureNode node) {
			return nodeTable[nodeTable.length - 1 - node.getOneIndex()];
		}

		/**
		 * Recursivly performs a pre-order traversal on a SuffixStructure of
		 * revered joined string and calculates failure links.
		 */
		private void findAllFailureLinks(SuffixStructureNode node, SuffixStructureNode lastValidNode) {
			if (lastValidNode!=null) {
				translateNode(node).setFailureLink(translateNode(lastValidNode));
			}
			if (!node.isRoot() && node.isSuffix()) {
				lastValidNode=node;
			}
			Iterator<SuffixStructureNode> iter = node.childrenIterator();
			while (iter.hasNext()) {
				SuffixStructureNode child = iter.next();
				findAllFailureLinks(child, lastValidNode);
			}
		}
		
		void findAllFailureLinks(StringSetProvider ssp) {
			SuffixStructure reversedSuffixStructure = ssp.getReverseSuffixStructure();
			// Log.getInstance().printInsane(SuffixStructureUtils.pstricksSuffixTree(reversedSuffixStructure));
			Log.getInstance().startTimer();			
			findAllFailureLinks(reversedSuffixStructure.getRoot(), null);
			Log.getInstance().stopTimer("calculate failure links");
		}
	}

	public int countNodes() {
		if (root==null) {
			return 0;
		} else {
			return root.countNodes();
		}
	}
	
	public void build(StringSetProvider ssp) {
		clear();
		Log.getInstance().startTimer();
		Log.getInstance().startTimer();
		// Table to maintain correspondence between string positions and node
		// objects
		AhoCorasickNode[] node_table = new AhoCorasickNode[ssp.getTotalLength()];
		// for (int i = 0; i < node_table.length; i++)	node_table[i] = null;
		Log.getInstance().restartTimer("init node table");

		List<String> l = ssp.getList();
		// step 1: build keyword trie
		root = new AhoCorasickNode(' ', 0);
		int nr = 1;
		// int pattern_nr=-1;
		int charIdx = 0;
		for (String s : l) {
			charIdx++;
			AhoCorasickNode node = root;
			for (int i = 0; i < s.length(); i++) {
				AhoCorasickNode child = node.getChildByChar(s.charAt(i));
				if (child == null) {
					child = new AhoCorasickNode(s.charAt(i), nr);
					nr++;
					node.addChild(child);
				}
				node = child;
				node_table[charIdx] = node;
				charIdx++;
			}
			if (node!=root) node.setOutput(node.getOutput()+1);
			node = root;
		}
		Log.getInstance().stopTimer("build keyword trie");

		FailureLinkFinder flf = new FailureLinkFinder(node_table);
		flf.findAllFailureLinks(ssp);
		
		// String s_forward = ssp.getJoinedString();
		// String s_rev = ssp.getReverseJoinedString();

		Log.getInstance().stopTimer("build aca");
	}

	public HMM createHMM(CharacterDistribution charDist) {
		ArrayList<Node> nodeList = new ArrayList<Node>();
		
		// Iterate over all nodes in pre-order and create one HMMNode for every
		// node visited.
		PreOrderIterator<AhoCorasickNode> iter = new PreOrderIterator<AhoCorasickNode>(root);
		while (iter.hasNext()) {
			AhoCorasickNode acNode = iter.next();

			// copy over array with character distribution and use it
			// to determine, which characters were already visited
			double[] charDistArray = charDist.getProbabilityArray(); 
			
			// create list of possible transitions by traversersing the 
			// list of failure links
			List<Transition> transitions = new ArrayList<Transition>(); 
			int output = 0;
			double p_rest = 1.0;
			AhoCorasickNode curNode = acNode;
			while (curNode!=null) {
				output+=curNode.getOutput();
				Iterator<AhoCorasickNode> childIter = curNode.children();
				// traverse children and check, if a character can be found that
				// was not yet used for a transition
				while (childIter.hasNext()) {
					AhoCorasickNode child = childIter.next();
					double p = charDistArray[child.getCharacter()];
					if (p>0.0) {
						transitions.add(new Transition(child.getNr(), p));
						charDistArray[child.getCharacter()]=0.0;
						p_rest-=p;
					}
				}
				// if this was root, we are finished
				if (curNode.getNr()==0) break;
				// if no failure link is present, check root
				if (curNode.getFailureLink()!=null) {
					curNode=curNode.getFailureLink();
				} else {
					curNode=root;
				}
			}
			// for those chars, for which no link has been found, transit
			// to root (with remaining probability)
			if (p_rest>0.0) {
				transitions.add(new Transition(0, p_rest));
			}
			// emission is deterministic in our case ...
			List<Emission> emissionDistribution = new ArrayList<Emission>();
			emissionDistribution.add(new Emission(output, 1.0));
			// create hmm node and add it to the list ...
			nodeList.add(new StandardNode(transitions, emissionDistribution));
		}
		return new HMM(nodeList);
	}

	public CDFA createCDFA(Alphabet alphabet) {
		List<int[]> transitionTable = new ArrayList<int[]>();
		List<Integer> outputTable = new ArrayList<Integer>();
		
		// Iterate over all nodes in pre-order and create one HMMNode for every
		// node visited.
		PreOrderIterator<AhoCorasickNode> iter = new PreOrderIterator<AhoCorasickNode>(root);
		while (iter.hasNext()) {
			AhoCorasickNode acNode = iter.next();
			
			// create list of possible transitions by traversersing the 
			// list of failure links
			int[] transitions = new int[alphabet.size()];
			Arrays.fill(transitions, -1);
			int output = 0;
			AhoCorasickNode curNode = acNode;
			while (curNode!=null) {
				output+=curNode.getOutput();
				Iterator<AhoCorasickNode> childIter = curNode.children();
				// traverse children and check, if a character can be found that
				// was not yet used for a transition
				while (childIter.hasNext()) {
					AhoCorasickNode child = childIter.next();
					int charIndex = alphabet.getIndex(child.getCharacter());
					if (charIndex<0) throw new IllegalStateException("Character not in alphabet");
					// if transition has already been defined for this character, continue
					if (transitions[charIndex]>=0) continue;
					// make entry for this transition
					transitions[charIndex]=child.getNr();
				}
				// if this was root, we are finished
				if (curNode.getNr()==0) break;
				// if no failure link is present, check root
				if (curNode.getFailureLink()!=null) {
					curNode=curNode.getFailureLink();
				} else {
					curNode=root;
				}
			}
			// for those chars, for which no link has been found, transit
			// to root
			for (int i=0; i<transitions.length; ++i) {
				if (transitions[i]<0) transitions[i]=0;
			}
			transitionTable.add(transitions);
			outputTable.add(output);
		}
		return new CDFA(alphabet, transitionTable, outputTable);
	}
	public AhoCorasickNode getRoot() {
		return root;
	}
	
	
}
